import cv2
import numpy as np

def load_digit_templates():
    """加载数字模板"""
    templates = []
    for i in range(10):
        template = cv2.imread(f'nums/{i}.png', cv2.IMREAD_GRAYSCALE)
        if template is not None:
            templates.append(template)
        else:
            print(f"警告: 无法加载模板 nums/{i}.png")
            templates.append(None)
    return templates

def preprocess_image(image_path):
    """预处理图像"""
    # 读取图片并转为灰度图
    gray = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if gray is None:
        raise ValueError(f"无法读取图像: {image_path}")
    
    # 使用与生成模板相同的方式进行二值化
    ret, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY|cv2.THRESH_OTSU)
    return gray, thresh

def find_digit_contours(thresh):
    """寻找数字轮廓"""
    contours = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)[0]
    
    # 将符合要求的轮廓收集并从左到右排序
    result = []
    for cnt in contours:
        [x, y, w, h] = cv2.boundingRect(cnt)
        # 根据数字大小过滤轮廓，确保是合适的数字尺寸
        if 22 > h > 10 and 16 > w > 8:  # 增加宽度过滤
            area = cv2.contourArea(cnt)
            # 添加面积过滤以确保是数字而不是噪声
            if area > 50:
                result.append([x, y, w, h])
    
    # 按x坐标排序（从左到右）
    result.sort(key=lambda x: x[0])
    return result

def recognize_digits(gray, thresh, contours, templates):
    """识别数字"""
    recognized_digits = []
    
    for i, (x, y, w, h) in enumerate(contours):
        # 在原图上画出轮廓
        cv2.rectangle(gray, (x, y), (x+w, y+h), (0, 0, 255), 1)
        
        # 提取数字区域并缩放到14×20大小
        digit_roi = cv2.resize(thresh[y:y+h, x:x+w], (14, 20))
        
        # 与模板匹配识别数字
        best_match = -1
        best_score = 0
        scores = []
        
        for digit, template in enumerate(templates):
            if template is not None:
                # 使用模板匹配
                match_result = cv2.matchTemplate(digit_roi, template, cv2.TM_CCORR_NORMED)
                # 确保提取单个值
                score_value = float(match_result[0, 0])
                scores.append((digit, score_value))
                if score_value > best_score:  # 不设置过高阈值，避免漏检
                    best_score = score_value
                    best_match = digit
        
        if best_match != -1 and best_score > 0.9:  # 稍微降低阈值
            recognized_digits.append(str(best_match))
            # 在图像上标注识别结果和置信度
            cv2.putText(gray, f"{best_match}({best_score:.2f})", (x, y-5), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 1)
        else:
            recognized_digits.append('?')  # 无法识别的数字
            # 显示所有可能的匹配结果（调试用）
            scores.sort(key=lambda x: x[1], reverse=True)
            top_matches = scores[:3]  # 显示前3个最佳匹配
            print(f"位置({x},{y})无法识别，最佳匹配: {top_matches}")
    
    return recognized_digits

def main():
    """主函数"""
    try:
        # 加载数字模板
        print("正在加载数字模板...")
        templates = load_digit_templates()
        
        # 预处理图像
        print("正在处理图像 123456789.jpg...")
        gray, thresh = preprocess_image('123456789.jpg')
        
        # 寻找数字轮廓
        print("正在寻找数字轮廓...")
        contours = find_digit_contours(thresh)
        print(f"找到 {len(contours)} 个数字轮廓")
        
        # 识别数字
        print("正在识别数字...")
        digits = recognize_digits(gray, thresh, contours, templates)
        
        # 输出结果
        result = ''.join(digits)
        print(f"识别结果: {result}")
        
        # 保存带标注的图像
        cv2.imwrite("digit_recognition_result.png", gray)
        print("结果已保存到 digit_recognition_result.png")
        
    except Exception as e:
        print(f"发生错误: {e}")

if __name__ == "__main__":
    main()